import Person;
import random;

from math import floor;

class Universe:
    '''A universe to store the network: the people and their connections'''
  
    def __init__(self, sim, n = 10, k = 3, p = 0.1, timezones = 0):
        '''Initialize the Universe with given number of people (n),
        given size of your close neighbourhood (k),
        and given probability of hooking up with a stranger (p).'''

        self.sim = sim
        self.persons = [];

        # SUGGESTION:
        #    Make close connection weights, and random connection weights
        #    configurable.
        # Q: Do we want to make them randomised in some sense?
        #   (MWW: Maybe not, to keep things simple)

        # ASSERTION: default time zone distribution is not having any
        if not timezones:
            for i in range(n):
                self.persons.append(Person.Person(i, sim, self, 0));
        elif timezones == 1:
            # ASSERTION: time zone distribution is evenly across
            # the four timezones in CONTUS.
            for i in range(n):
                timezone = int(floor(random.random()*4));
                self.persons.append(Person.Person(i, sim, self, timezone));

        # elif:
            # split weirdly across the globe.
            # not yet sure how to do it, though
                
        # self.connections = [ [ 0 for i in range(n)] for j in range(n) ]
        self.n = n;
        self.k = k;
        self.p = p;
        self.graph_name = "";

        self.newman_watts_strogatz_graph(self.k, self.p);

    def connection_set(self, person1, person2, value):
        '''Connect/disconnect person1 and person2'''

        p1_index = self.persons.index(person1)
        p2_index = self.persons.index(person2)
        # self.connections[p1_index][p2_index] = value


    def connection_get(self, person1, person2):
        '''Return the connection person1 and person2'''

        p1_index = self.persons.index(person1)
        p2_index = self.persons.index(person2)
        # return self.connections[p1_index][p2_index]


#    def connections_show(self):
#        '''Show all connections'''

#        for person1 in self.persons:
#            for person2 in self.persons:
                # print "%d  " % self.connection_get(person1, person2),
                # print

    # Modified Newman-Watts-Strogatz small world graph generator
    # taken from https://networkx.lanl.gov/trac/browser/networkx/trunk/networkx/generators/random_graphs.py
    def newman_watts_strogatz_graph(self, k, p, seed=None):
        """Return a Newman-Watts-Strogatz small world graph.
    
        Parameters
        ----------
        k : int
        Each node is connected to k nearest neighbors in ring topology
        p : float
        The probability of adding a new edge for each edge
        seed : int, optional       
        seed for random number generator (default=None)
        
        Notes
        -----
        First create a ring over n nodes.  Then each node in the ring is
        connected with its k nearest neighbors (k-1 neighbors if k is odd). 
        Then shortcuts are created by adding new edges as follows:
        for each edge u-v in the underlying "n-ring with k nearest neighbors"
        with probability p add a new edge u-w with randomly-chosen existing
        node w.  In contrast with watts_strogatz_graph(), no edges are removed.
        
        See Also
        --------
        watts_strogatz_graph()
        
        References
        ----------
        .. [1] M. E. J. Newman and D. J. Watts,
        Renormalization group analysis of the small-world network model,
        Physics Letters A, 263, 341, 1999.
        http://dx.doi.org/10.1016/S0375-9601(99)00757-4
        """

        # QUESTION: do we want to switch to a Watts-Strogatz graph, where
        #   the random connections can remove an existing close world connection?
        # MWW: I'm suspicious that it might not be very useful, but it was a thought,
        #      so I figured I'd add it.

        if seed is not None:
            random.seed(seed);

        self.graph_name="newman_watts_strogatz_graph(%s,%s,%s)" % (self.n,k,p);

        nlist = [];
        edges = [];
        for i in range(self.n):
            nlist.append(i);
        fromv = nlist;
        # connect the k/2 neighbors
        for n in range(1, k/2+1):
            tov = fromv[n:] + fromv[0:n] # the first n are now last
            for i in range(len(fromv)):
                # connect the two nodes
                u = fromv[i];
                v = tov[i];
                # we reciprocally follow each other
                edges.append((u, v));
                edges.append((v, u));
                # self.connections[u][v] = 1;
                # self.connections[v][u] = 1;
                # added list of people you follow, for easier search
                if v not in self.persons[u].follows:
                    self.persons[u].follows.append((v, 1));
                if u not in self.persons[v].follows:
                    self.persons[v].follows.append((u, 1));

        # for each edge u-v, with probability p, randomly select existing
        # node w and add new edge u-w
        for (u, v) in edges:
            if random.random() < p:
                w = random.choice(nlist);
                # no self-loops and reject if edge u-w exists
                # is that the correct NWS model?
                while (w == u) or (w in self.persons[u].follows):
                    w = random.choice(nlist);
                # self.connections[u][w] = 1;
                # added list of people you follow, for easier search
                self.persons[u].follows.append((w, 0.1));
        # self.connections_show();                
        return;
